#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/3/6 15:14
# @Author  : wanghaoran
# @File    : query_transform.py
from advanced_module.base import BasePromptEngineTool
from llm_models import chat_qwen_72b


class QueryTransform(BasePromptEngineTool):
    def __init__(self):
        self.prompt = \
        """
            针对以下问题，判断其下一步需要进行什么操作，并以列表形式返回其对应的序号（1、2、3、4）和操作后得到的结果（返回示例：[1, ["langchain在github上的star数多少？"， "llamaindex在github上的star数多少？"]]）：
            1.问题高度口语化、不标准，需要进行重写，以列表形式返回重写问题；
            2.问题复杂，能够被分解为多个更简单的子问题，各个子问题之间能够独立进行回答，例如当你需要比较多个事物的时候，以列表形式返回分解得到的子问题；
            3.问题复杂，需要多方面的信息，并且单个查询不足以完全理解或检索必要的信息，问题需要分解为递进的子问题（下一个问题的构建需要前一个问题的回答），例如在信息分层且需要迭代探索的环境中，以列表形式返回分解得到的子问题，其中需要检索答案的地方以##进行填充；
            4.问题描述直接、准确，能够通过单次检索获取所需信息，不需要进行复杂的转换操作，例如直接询问某个事物的信息，这种情况不需要任何返回值，返回[4, ""]即可。
            ## 问题
            {query}
            ## 返回结果
        """

    def get_response(self, query):
        prompt = self.prompt.format(query=query)
        response = chat_qwen_72b(prompt)
        try:
            response = eval(response)
            if response[0] == 4:
                return [query]
            return response[1]
        except:
            return [query]


class QueryExpansion(BasePromptEngineTool):
    def __init__(self):
        self.prompt = \
        """
            你是一个问答助手，你的用户正在向你询问问题。针对用户的问题，你需要提出至多5个（可以少于5个）相关问题来帮助他们找到他们所需要的信息。
            提出的问题需要满足以下条件：问题简短、不是复合句、需要涵盖不同方向、确保是完整的问题、问题之间无交叉、与原问题高度相关。
            每行输出一个问题，不要给问题生成编号。
            例：
            ## 问题：酒驾应该如何处罚？
            ## 扩展问题：
            酒驾的法律定义是什么？
            酒驾的处罚标准是什么？
            酒驾对驾驶执照有什么影响？
            酒驾会导致哪些刑事责任？
            ## 问题：{query}
            ## 扩展问题：
        """

    def get_response(self, query):
        prompt = self.prompt.format(query=query)
        response = chat_qwen_72b(prompt)
        expand_query_list = response.split("\n")
        expand_query_list.append(query)
        return expand_query_list


if __name__ == "__main__":
    query_transformer = QueryTransform()
    # “重庆科技园”的基本信息？
    # 执导电影《异形4》的导演今年的年龄是多少？
    res = query_transformer.get_response("“重庆科技园”的基本信息？")
    print(res)
